package org.jflame.context.ratelimiter;

import org.jflame.context.cache.redis.RedisClient;

import java.util.Arrays;
import java.util.List;

public class RedisRateLimiter implements RateLimiter {

    private RedisClient redisClient;
    private String luaScript;

    public RedisRateLimiter() {
        luaScript = new StringBuilder().append("local key1 = KEYS[1] ")
                .append(" local expire = tonumber(ARGV[1])")
                .append(" local times = tonumber(ARGV[2])")
                .append(" local val = redis.call('incr', key1)")
                .append(" if val == 1 then ")
                .append("    redis.call('expire', key1, expire)")
                .append(" else ")
                .append("   local ttl = redis.call('ttl', key1)")
                .append("   if ttl == -1 then")
                .append("     redis.call('expire', key1, expire)")
                .append("    end")
                .append(" end")
                .append(" if val > times then")
                .append("  return 0")
                .append(" end ")
                .append(" return 1")
                .toString();

    }

    @Override
    public boolean checkLimit(String requestId, int threshold, long intervalInSecond) {
        List<Long> args = Arrays.asList(intervalInSecond, (long) threshold);
        Long r = redisClient.runScript(luaScript, Arrays.asList(requestId), args, Long.class);
        return r == 0;
    }

    public void setRedisClient(RedisClient redisClient) {
        this.redisClient = redisClient;
    }

}
